use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{parse_macro_input, FnArg, ItemTrait, ReturnType, TraitItem};

#[proc_macro_attribute]
pub fn generate_proxy(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as ItemTrait);
    let ident = &input.ident;
    let visibility = &input.vis;

    let proxy_ident = format_ident!("{ident}Proxy");

    let mut sync_methods = vec![];
    let mut async_methods = vec![];

    for item in &input.items {
        // Need to generate wrappers for each fn.
        if let TraitItem::Fn(method) = item {
            let sig = &method.sig;
            let name = &sig.ident;

            // Extract the relevant parts of the signature.
            let mut self_arg = None;
            let mut other_inputs = vec![];
            let mut arg_names = vec![];

            for input in &sig.inputs {
                match input {
                    FnArg::Receiver(rec) => {
                        self_arg = Some(quote! { #rec });
                    }
                    FnArg::Typed(pat_type) => {
                        arg_names.push(match &*pat_type.pat {
                            syn::Pat::Ident(pat_ident) => quote! { #pat_ident },
                            _ => panic!("Unsupported pattern in argument"),
                        });
                        other_inputs.push(FnArg::Typed(pat_type.clone()));
                    }
                }
            }

            // Check that there really is a self arg there.
            if self_arg.is_none() {
                self_arg.expect("Expected a self receiver in trait method");
            }

            // Extract the return type.  This is kinda a hack if there is no return type.
            let ret_type = match &sig.output {
                ReturnType::Type(_, ty) => quote! { #ty },
                ReturnType::Default => quote! { () },
            };

            // Simple dynamic dispatch.
            let blocking_name = format_ident!("{}_blocking", name);
            sync_methods.push(quote! {
                pub fn #blocking_name(&self, #(#other_inputs),*) -> #ret_type {
                    self.inner.#name(#(#arg_names),*)
                }
            });

            // Here we use spawn_blocking to make it safe.
            let async_name = format_ident!("{}_async", name);
            async_methods.push(quote! {
                pub async fn #async_name(&self, #(#other_inputs),*) -> #ret_type
                where
                    #ret_type: Send + 'static,
                    Self: Sync + Send + 'static,
                {
                    // FIXME should we handle panics better?
                    // FIXME shouldwe get rid of hard tokio dep here?
                    let inner = self.inner.clone();
                    tokio::task::spawn_blocking(move || {
                        inner.#name(#(#arg_names),*)
                    })
                    .await
                    .expect("proxy: worker thread panicked")
                }
            });
        }
    }

    let expanded = quote! {
        #input

        #visibility struct #proxy_ident {
            inner: std::sync::Arc<dyn #ident + Send + Sync>,
        }

        impl #proxy_ident {
            pub fn new<T: #ident + Send + Sync + 'static>(t: T) -> Self {
                Self {
                    inner: std::sync::Arc::new(t),
                }
            }

            #(#sync_methods)*

            #(#async_methods)*
        }
    };

    TokenStream::from(expanded)
}
